"""
python fujian1_SplineAndArima_gragh.py
"""

import os
import json
import pandas as pd
import matplotlib.pyplot as plt

# 设置文件路径
spline_dir = r"./fujian/fujian1/cubic_spline/interpolation_output"
# pmdarima_dir = r"./fujian/fujian1/pmdarima"
pmdarima_dir = r"./fujian/fujian1/pmdarima_single_batch"

output_plot_dir = os.path.join(pmdarima_dir, "comparison_plots")

# 创建保存图像的文件夹（如果不存在）
os.makedirs(output_plot_dir, exist_ok=True)

# 获取两个目录中的 JSON 文件列表
spline_files = [f for f in os.listdir(spline_dir) if f.endswith('.json')]
pmdarima_files = [f for f in os.listdir(pmdarima_dir) if f.endswith('.json')]

# 遍历 spline 目录中的每个文件
for spline_file in spline_files:
    # 获取相同名字的 pmdarima 文件路径
    # pmdarima_file = os.path.join(pmdarima_dir, spline_file)
    # pmdarima_file = os.path.join(pmdarima_dir, spline_file)
    
    category_number = spline_file.split('_')[1].split('.')[0]  # 获取 'category348' 中的 '348'
    
    # 生成对应的 pmdarima 文件名
    pmdarima_file_name = f"smoothAfterSpline_{category_number}.json"
    print("pmdarima file name:", pmdarima_file_name)
    pmdarima_file = os.path.join(pmdarima_dir, pmdarima_file_name)
    
    
    # 检查 pmdarima 文件是否存在
    if os.path.exists(pmdarima_file):
        spline_file_path = os.path.join(spline_dir, spline_file)

        # 读取样条插值数据
        with open(spline_file_path, 'r') as f:
            spline_data = json.load(f)
        
        # 读取平滑+预测数据
        with open(pmdarima_file, 'r') as f:
            pmdarima_data = json.load(f)
        
        # 转换为 DataFrame 并处理日期格式
        spline_df = pd.DataFrame(spline_data)
        pmdarima_df = pd.DataFrame(pmdarima_data)

        # 将日期转换为时间格式以便按月排序
        spline_df['date'] = pd.to_datetime(spline_df['date'])
        pmdarima_df['date'] = pd.to_datetime(pmdarima_df['date'])

        # 设置绘图
        plt.figure(figsize=(10, 6))
        plt.plot(spline_df['date'], spline_df['inventory'], label="Spline Interpolation", color="blue", marker='o')
        plt.plot(pmdarima_df['date'], pmdarima_df['inventory'], label="Smooth + Forecast", color="orange", marker='x')

        # 图形细节设置
        plt.title(f"Category {spline_file.split('_')[-1].replace('.json', '')}")
        plt.xlabel("Date (YYYY-MM)")
        plt.ylabel("Inventory")
        plt.legend()
        plt.grid(True)
        
        # 日期显示格式
        plt.gca().xaxis.set_major_formatter(plt.matplotlib.dates.DateFormatter('%Y-%m'))

        # 保存图像
        output_path = os.path.join(output_plot_dir, f"{spline_file.replace('.json', '')}_comparison_plot.png")
        plt.savefig(output_path, bbox_inches='tight')
        plt.close()

        print(f"Saved plot for {spline_file} at {output_path}")
    else:
        print(f"Skipping {spline_file}, no matching PMDARIMA file found.")
